/*
 * Copyright (c) Extreme Vision Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "yoloworld_prepostprocessor.h"
namespace ev
{
    namespace vision
    {
        float YOLOWorldPreprocessor::GetScale() { return m_scale; }

        static float sigmoid(const float x) { return 1.f / (1.f + exp(-x)); }

        EVStatus YOLOWorldPreprocessor::Run(cv::Mat &in_mat, cv::Mat &out_mat, int out_width)
        {
            if (in_mat.empty())
            {
                return EVStatus::EV_FAIL;
            }
            REC_TIME(t0);
            int src_h = in_mat.rows;
            int src_w = in_mat.cols;
            int max_size = src_h > src_w ? src_h : src_w;
            m_scale = out_width * 1.0f / max_size;
            int pad_h = (max_size - src_h) / 2;
            int pad_w = (max_size - src_w) / 2;
            cv::Mat pad_img = cv::Mat::zeros(max_size, max_size, CV_8UC3);
            cv::Rect rect{pad_w, pad_h, src_w, src_h};
            in_mat.copyTo(pad_img(rect));
            cv::Mat board, resized_mat;
            cv::resize(pad_img, resized_mat, cv::Size{out_width, out_width});
            resized_mat.convertTo(board, CV_32FC3, 1 / 255.);
            out_mat = board.clone();
            std::vector<cv::Mat> chw_wrappers;
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1,
                                      out_mat.data + 2 * sizeof(float) * out_width * out_width);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1,
                                      out_mat.data + sizeof(float) * out_width * out_width);
            cv::split(board, chw_wrappers);
            REC_TIME(t1);
            EVLOG(INFO) << "YOLOWorldPreprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EVStatus::EV_SUCCESS;
        }

        EVStatus YOLOWorldPostprocessor::Run(EVMatData *out, std::vector<BoxInfo> &objects, float scale, float thresh,
                                             const int img_w, const int img_h, float nms_thresh, bool class_aware)
        {
            REC_TIME(t0);
            int channel = out->dims[2];
            int box_num = out->dims[1];
            int class_num = channel - 4;
            float *data = (float *)out->data;
            int net_size = std::max<int>(img_w, img_h) * scale;
            int pad_h = (static_cast<int>(net_size / scale) - img_h) / 2;
            int pad_w = (static_cast<int>(net_size / scale) - img_w) / 2;
            for (int i = 0; i < box_num; ++i)
            {
                float max_score = data[4];
                int index = 0;
                for (int j = 5; j < channel; j++)
                {
                    if (data[j] > max_score)
                    {
                        max_score = data[j];
                        index = j - 4;
                    }
                }
                if (max_score > thresh)
                {
                    float x1 = data[0] / scale - pad_w;
                    float y1 = data[1] / scale - pad_h;
                    float x2 = data[2] / scale - pad_w;
                    float y2 = data[3] / scale - pad_h;
                    BoxInfo box_info = {x1, y1, x2, y2, max_score, index};
                    box_info.x1 = std::max(box_info.x1, 0.0f);
                    box_info.x2 = std::min(box_info.x2, img_w - 1.0f);
                    box_info.y1 = std::max(box_info.y1, 0.0f);
                    box_info.y2 = std::min(box_info.y2, img_h - 1.0f);
                    objects.push_back(box_info);
                }
                data = data + channel;
            }
            NMS(objects, nms_thresh, class_aware);
            REC_TIME(t1);
            EVLOG(INFO) << "YOLOWorldPreprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }
    } // namespace vision
} // namespace ev
